package com.zhengb.flink.study.datastream.sourcefunction;

import lombok.extern.slf4j.Slf4j;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.streaming.api.functions.source.ParallelSourceFunction;

/**
 * 自定义并行 source function
 *
 * @author zhengbo
 * @date 2019/12/19
 */
@Slf4j
public class CustomParallelSourceFunction {

    public static class CustomSourceFunction implements ParallelSourceFunction<Long> {

        private long count = 1L;

        private static volatile boolean isRunning = true;

        @Override
        public void run(SourceContext<Long> ctx) throws Exception {

            while (isRunning) {

                log.info("generate number:{}", count);

                ctx.collectWithTimestamp(count, System.currentTimeMillis());
                count++;
                Thread.sleep(1000);
            }
        }

        @Override
        public void cancel() {
            isRunning = false;
        }
    }

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.addSource(new CustomSourceFunction()).setParallelism(5).addSink(new SinkFunction<Long>() {
            @Override
            public void invoke(Long value, Context context) throws Exception {
                log.info("sink data:{}", value);
            }
        });

        env.execute("CustomParallelSourceFunction");
    }


}
